"""
BaoStock日K数据DAO层

本文件提供BaoStock日K数据的数据库操作。
相关知识点：
- SQLAlchemy异步ORM
- 数据库操作
"""
from datetime import date
from typing import Sequence, List
import pandas as pd
from sqlalchemy import delete, desc, insert, select, update, distinct
from sqlalchemy.ext.asyncio import AsyncSession
from module_stock.entity.do.baostock_daily_do import StockBaoStockDaily
from utils.page_util import PageUtil
from utils.pd_util import AsyncORMConverter


class BaoStockDailyDao:
    """
    BaoStock股票日线模块数据库操作层
    """

    @classmethod
    async def get_baostock_daily_detail_by_id(cls, db: AsyncSession, id: int):
        """
        根据id获取BaoStock股票日线详细信息

        :param db: orm对象
        :param id: id
        :return: BaoStock股票日线信息对象
        """
        daily_info = (
            (
                await db.execute(
                    select(StockBaoStockDaily)
                    .where(StockBaoStockDaily.id == id)
                )
            )
            .scalars()
            .first()
        )
        return daily_info

    @classmethod
    async def get_last_baostock_daily_detail(cls, db: AsyncSession, symbol: str):
        """
        根据股票代码获取数据表最后一个日线信息

        :param db: orm对象
        :param symbol: 股票代码
        :return: BaoStock股票日线信息对象
        """
        daily_info = (
            (
                await db.execute(
                    select(StockBaoStockDaily)
                    .where(StockBaoStockDaily.symbol == symbol)
                    .order_by(desc(StockBaoStockDaily.trade_date))
                    .limit(1)
                )
            )
            .scalars()
            .first()
        )
        return daily_info

    @classmethod
    async def get_all_baostock_daily(cls, db: AsyncSession, symbol: str):
        """
        获取所有BaoStock股票日线信息

        :param db: orm对象
        :param symbol: 股票代码
        :return: BaoStock股票日线信息集
        """
        query = select(StockBaoStockDaily).where(
            StockBaoStockDaily.symbol == symbol
        ).order_by(StockBaoStockDaily.trade_date.asc())
        
        daily_list = (await db.execute(query)).scalars().all()
        return daily_list

    @classmethod
    async def get_all_baostock_daily_after_date(
        cls, db: AsyncSession, symbol: str, start_date: date
    ) -> Sequence[StockBaoStockDaily]:
        """
        获取指定日期之后的BaoStock股票日线信息

        :param db: orm对象
        :param symbol: 股票代码
        :param start_date: 起始日期
        :return: BaoStock股票日线信息集
        """
        query = (
            select(StockBaoStockDaily)
            .where(
                StockBaoStockDaily.symbol == symbol,
                StockBaoStockDaily.trade_date >= start_date
            )
            .order_by(StockBaoStockDaily.trade_date.asc())
        )
        daily_list = (await db.execute(query)).scalars().all()
        return daily_list

    @classmethod
    async def get_all_baostock_daily_pd(cls, db: AsyncSession, symbol: str):
        """
        获取所有BaoStock股票日线信息（返回DataFrame）

        :param db: orm对象
        :param symbol: 股票代码
        :return: DataFrame
        """
        query = (
            select(StockBaoStockDaily)
            .where(StockBaoStockDaily.symbol == symbol)
            .order_by(StockBaoStockDaily.trade_date.asc())
        )
        dailys = (await db.execute(query)).scalars().all()
        daily_df = pd.DataFrame(AsyncORMConverter._orms_to_dicts(dailys))
        return daily_df

    @classmethod
    async def get_existing_dates(
        cls, db: AsyncSession, symbol: str, trade_dates: List[date]
    ) -> List[date]:
        """
        获取已存在的交易日期列表

        :param db: orm对象
        :param symbol: 股票代码
        :param trade_dates: 交易日期列表
        :return: 已存在的交易日期列表
        """
        if not trade_dates:
            return []
        
        query = select(StockBaoStockDaily.trade_date).where(
            StockBaoStockDaily.symbol == symbol,
            StockBaoStockDaily.trade_date.in_(trade_dates)
        )
        result = await db.execute(query)
        existing_dates = list(result.scalars().all())
        return existing_dates

    @staticmethod
    async def get_all_symbols(query_db: AsyncSession) -> List[str]:
        """
        获取已保存的BaoStock日K数据的所有股票代码
        
        :param query_db: 数据库会话
        :return: 股票代码列表
        """
        stmt = select(distinct(StockBaoStockDaily.symbol)).order_by(StockBaoStockDaily.symbol)
        result = await query_db.execute(stmt)
        return [symbol for symbol, in result.all()]

    @classmethod
    async def batch_add_baostock_daily_dao(cls, db: AsyncSession, dailys: List[dict]):
        """
        批量新增BaoStock股票日线数据库操作

        :param db: orm对象
        :param dailys: BaoStock股票日线实体对象列表
        :return:
        """
        if not dailys:
            return
        
        st = insert(StockBaoStockDaily).values(dailys)
        await db.execute(st)

    @classmethod
    async def batch_delete_baostock_daily_dao(cls, db: AsyncSession, symbols: List[str]):
        """
        批量删除BaoStock股票日线数据库操作

        :param db: orm对象
        :param symbols: 股票代码列表
        :return:
        """
        if not symbols:
            return
        
        await db.execute(
            delete(StockBaoStockDaily).where(StockBaoStockDaily.symbol.in_(symbols))
        )

    @classmethod
    async def delete_baostock_daily_by_symbol_and_date(
        cls, db: AsyncSession, symbol: str, trade_date: date
    ):
        """
        根据股票代码和交易日期删除BaoStock股票日线数据

        :param db: orm对象
        :param symbol: 股票代码
        :param trade_date: 交易日期
        :return:
        """
        await db.execute(
            delete(StockBaoStockDaily).where(
                StockBaoStockDaily.symbol == symbol,
                StockBaoStockDaily.trade_date == trade_date
            )
        )

